import cv2
import numpy as np

class yolo_detect:

    def __init__(self, params, target, draw=False):

        """
        @ param : confThreshold : Confidence threshold
        @ param : nmsThreshold : Non-maximum suppression threshold
        @ param : inpWidth : Width of network's input image
        @ param : inpHeight : Height of network's input image
        @ param : classesFile : Load names of classes
        @ param : modelConfiguration : Give the configuration files for the model
        @ param : modelWeights : Give the weight files for the model
        """
        
        self.confThreshold = 0.5
        self.nmsThreshold = 0.65
        self.inpWidth = 416
        self.inpHeight = 416

        # Load names of classes
        self.classesFile = params['class_names']
        self.classes = None
        self.classes_dict = {}
        self.target = target
        with open(self.classesFile, 'rt') as f:
            self.classes = f.read().rstrip('\n').split('\n')

        for i in range(0, len(self.classes)):
            self.classes_dict[self.classes[i]] = i

        # Give the configuration and weight files for the model and load the network using them.
        self.modelConfiguration = params['net_cfg']
        self.modelWeights = params['weight_cfg']

        self.net = cv2.dnn.readNetFromDarknet(self.modelConfiguration, self.modelWeights)
        self.net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
        self.net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)

        self.draw = draw

    def getOutputsNames(self):

        """
        Get the names of the output layers
        """

        # Get the names of all the layers in the network
        layersNames = self.net.getLayerNames()

        # Get the names of the output layers, i.e. the layers with unconnected outputs
        return [layersNames[i[0] - 1] for i in self.net.getUnconnectedOutLayers()]

    def drawPred(self, frame, classId, conf, left, top, right, bottom):

        """
        Draw the predicted bounding box
        """

        # Draw a bounding box.
        cv2.rectangle(frame, (left, top), (right, bottom), (255, 178, 50), 3)
        
        label = '%.2f' % conf
            
        # Get the label for the class name and its confidence
        if self.classes:
            assert(classId < len(self.classes))
            label = '%s:%s' % (self.classes[classId], label)

        #Display the label at the top of the bounding box
        #labelSize, baseLine = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        #top = max(top, labelSize[1])
        #cv.rectangle(frame, (left, top - round(1.5*labelSize[1])), (left + round(1.5*labelSize[0]), top + baseLine), (255, 255, 255), cv.FILLED)
        #cv.putText(frame, label, (left, top), cv.FONT_HERSHEY_SIMPLEX, 0.75, (0,0,0), 1)

    def postprocess(self, frame, outs):

        """
        Remove the bounding boxes with low confidence using non-maxima suppression

        @ param frame : the input frame
        @ param outs : the output of the yolo net
        @ return : the bbox list
        """

        frameHeight = frame.shape[0]
        frameWidth = frame.shape[1]

        # Scan through all the bounding boxes output from the network and keep only the
        # ones with high confidence scores. Assign the box's class label as the class with the highest score.
        classIds = []
        confidences = []
        boxes = []
        bbox_for_tracking = []
        for out in outs:
            for detection in out:
                scores = detection[5:]
                classId = np.argmax(scores)
                confidence = scores[classId]
                if confidence > self.confThreshold:
                    center_x = int(detection[0] * frameWidth)
                    center_y = int(detection[1] * frameHeight)
                    width = int(detection[2] * frameWidth)
                    height = int(detection[3] * frameHeight)
                    left = int(center_x - width / 2)
                    top = int(center_y - height / 2)
                    classIds.append(classId)
                    confidences.append(float(confidence))
                    boxes.append([left, top, width, height])
                    
        # Perform non maximum suppression to eliminate redundant overlapping boxes with
        # lower confidences.
        indices = cv2.dnn.NMSBoxes(boxes, confidences, self.confThreshold, self.nmsThreshold)
        for i in indices:
            i = i[0]
            box = boxes[i]
            left = box[0]
            top = box[1]
            width = box[2]
            height = box[3]
            for t in self.target:
                if classIds[i] == self.classes_dict[t]:
                    bbox_for_tracking.append((left, top, width, height))
                    if self.draw == True:
                        self.drawPred(frame, classIds[i], confidences[i], left, top, left + width, top + height)

        return bbox_for_tracking

    def yolo_pipeline(self, frame):
        """
        Create a yolo net and run the yolo

        @ param frame : the input frame
        @ return : the list of the bbox
        """
        net = self.net
        # Create a 4D blob from a frame.
        blob = cv2.dnn.blobFromImage(frame, 1/255, (self.inpWidth, self.inpHeight), [0,0,0], 1, crop=False)

        # Sets the input to the network
        net.setInput(blob)

        # Runs the forward pass to get output of the output layers
        outs = net.forward(self.getOutputsNames())

        # Remove the bounding boxes with low confidence
        bbox = self.postprocess(frame, outs)
        max_area = 0
        max_box = None
        for subbox in bbox:
            if max_area < subbox[2] * subbox[3]:
                max_area = subbox[2] * subbox[3]
                max_box = subbox
        return [max_box]
